#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/7/18 21:11
# @Email   : you.ze.shun@qq.com
# @File    : _load.py
# https://python.langchain.com/docs/modules/data_connection/vectorstores/
import os
import getpass
from langchain.document_loaders import TextLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import FAISS

# 会显示提示字符串, 关闭键盘的屏幕回显，然后读取密码
os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key:')
os.environ['OPENAI_API_BASE'] = getpass.getpass('OPENAI_API_BASE:')

raw_documents = TextLoader('_similarity_search.txt').load()
text_splitter = CharacterTextSplitter()
documents = text_splitter.split_documents(raw_documents)

embeddings = OpenAIEmbeddings()
db = FAISS.from_documents(documents, embeddings)
query = "学历高才有光明的未来"
docs = db.similarity_search(query, k=1)
print(docs)
for i, v in enumerate(docs):
    print(f'paragraph {i}:{v.page_content}')
embedding_vector = embeddings.embed_query(query)
docs = db.similarity_search_by_vector(embedding_vector, k=1)
print(docs)
